from langchain_core.tools import tool
from app.bailian.common import chat_prompt_template, llm
from pydantic import BaseModel, Field


"绑定自定义工具-使用tool装饰器"


class FetchWeatherInputArgs(BaseModel):
    city: str = Field(description="城市名称")


# 1.开发自定义工具, 并注册tool
@tool(
    description="获取某个城市的真实天气",
    args_schema=FetchWeatherInputArgs
)
def fetch_weather(city: str) -> str:
    """获取某个城市的真实天气"""
    # 模拟调用天气API逻辑...
    weather_data = {
        '北京': '多云',
        '深圳': '晴朗'
    }
    if city not in weather_data.keys():
        return '天气多变'
    return weather_data[city]


# 定义工具字典, value为python方法
tool_dict = {
    'fetch_weather': fetch_weather
}
# 3.将大模型与Tool对象绑定
llm_with_tool = llm.bind_tools([fetch_weather])
# 链式声明
chain = chat_prompt_template | llm_with_tool
# 4.调用大模型
resp = chain.invoke(input={"role": "气象学", "domain": "天气领域", "question": "深圳天气怎么样?"})
# 5.调用工具
for tool_calls in resp.tool_calls:
    # 获取要调用的参数和工具
    args = tool_calls["args"]
    func_name = tool_calls["name"]
    tool_func = tool_dict[func_name]
    # 调用工具
    tool_content = tool_func.invoke(args)
    print(tool_content)  # 晴朗
